"""
  描述：
        后方HDV处于跟随状态
        hdvH(hdv head)表示头部自由驾驶的hdv
        hdvF(hdv follow)表示尾部跟随状态的hdv
        cavM表示换道CAV车辆
        cav1、cav2、cav3...表示队列中的cav

  迭代计划：
        v1.0.0: 搭建hdvH和cavM的简单交互Env场景、搭建SAC算法、搭建测试流程、绘制测试数据。
        v2.0.0: 搭建hdvH、cav1、hdvF车队和cavM的交互Env场景, cavM换道插入到hdvH和cav1之间。 
                需要引入transformer模型去生成cavM的真实控制数据。
        v3.0.0: 开始考虑传感器测量的不确定性, 即加入卡尔曼滤波去处理测量值。
            目前的问题：
                1. todo 需要解决 感觉bicyle_model的卡尔曼滤波器有问题, 模型计算的预测值和环境迭代的真实值偏差比较大
                   (已解决, 因为bicyle_model和kalman滤波的航向角计算公式有偏差)
        v4.0.0: 开始考虑hdvH驾驶行为的不确定性, 即引入transformer对车辆的未来驾驶行为进行预测。
                1. 除了用transformer去预测hdvH和hdvF, 让RL控制器也生成2帧的预测控制。
                2. 将预测状态引入到reward函数中(这种带有预测的方式, 在实验时貌似看到的不错的)
        v5.0.0: 开始考虑transformer模型预测的不确定性
                1. 通过使用一个模型在训练和预测的时候都开启mcDropout来形成多种预测数据。
                2. todo 需要实现 需要将transformer的预测方差引入到reward函数中。
        v6.0.0: 在控制器生成的控制量中加入控制扰动(用来模拟控制的不确定性)
        v7.0.0: todo
"""

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

# 导入外部包
import numpy as np
import random
import math
import copy
from sklearn.model_selection import train_test_split
from keras.models import load_model

# 导入内部包
from Dynamic_Model.point_mass_model import PointMassModel
from Dynamic_Model.bicyle_model import BicycleModel
from rl_controller.Env.Follow_Env_Reward import *
from utils.transformer import PositionalEncoding
from utils.lstm import normalization, anit_normalization, create_lstm_dataset

# 设置种子
random.seed(100)

# 常量
BASE_ACTION_BOUND = 1 # 默认动作范围
ACC_BOUND = 3 # 加速度范围
DELTA_BOUND = 0.43633 # 转向角范围
N_ACTION = 9 # todo 动作空间维度（M的一帧轴向加速度和转向角、M的两帧预测加速度和转向角、cav1的一帧纵向加速度、cav1的两帧预测纵向加速度）
N_STATE = 23 # todo 状态空间维度（M和hdvH的纵向间距、M的横向换道累计距离、M的轴向速度、M的航向角、hdvH的纵向速度、M和cav1的纵向间距、cav1和hdvF的纵向间距、cav1的纵向速度、hdvF的纵向速度、hdvH的加速度、hdvF的加速度, cavM的y协方差、cavM的x协方差、cavM的v协方差、cavM的航向角协方差, hdvH的两帧预测加速度, hdvF的两帧预测加速度, hdvH的两帧预测方差, hdvF的两帧预测方差）
FREE_N_PAST = 30 # 头车的历史帧数（过去 + 当前）
FREE_N_FUTURE = 10 # 预测头车未来状态的帧数
FOLLOW_N_PAST = 30 # 后方HDV的历史帧数（过去 + 当前）
FOLLOW_N_FUTURE = 10 # 预测HDV未来状态的帧数
TCAV = 1.3 # 与CAV车辆的期望车头时距
THDV = 1.8 # 与HDV车辆的期望车头时距
MAX_ACC_JERK = 40 # 最大加加速度
ACC_NOISE_STD = 0.05 # 加速度执行时的扰动
DELTA_NOISE_STD = 0.05 # 转向角执行时的扰动

TRANSFORMER_FOLLOW_MODEL_TRUE = "../../transformer/model/follow_EPOCHS_50_BATCH_SIZE_64_transformer-normalization.keras" # 生成hdvF真实的控制行为
TRANSFORMER_FOLLOW_MODEL_PRED = "../../transformer/model/follow_EPOCHS_50_BATCH_SIZE_64_transformer-normalization.keras" # todo 换模型，暂时没训练，就用同一个 生成hdvF预测的控制行为
TRANSFORMER_FREE_DRIVING_MODEL = "../../transformer/model/free_drive_EPOCHS_50_BATCH_SIZE_64_transformer-normalization.keras" # 生成hdvH的预测行为

class FollowEnv:
  def __init__ (self):
    # 状态参数
    self.dt = 0.1 # 每次状态跟新的时间间隔
    self.lane_width = 3.75 # 车道宽度为3.75m
    self.hdvH_n_past = FREE_N_PAST
    self.hdvH_n_future = FREE_N_FUTURE
    self.hdvF_n_past = FOLLOW_N_PAST
    self.hdvF_n_future = FOLLOW_N_FUTURE

    # 维度参数
    self.action_bound = BASE_ACTION_BOUND
    self.acc_bound = ACC_BOUND
    self.delta_bound = DELTA_BOUND
    self.n_action = N_ACTION
    self.n_state = N_STATE

    # reward参数
    self.tcav = TCAV
    self.thdv = THDV
    self.max_acc_jerk = MAX_ACC_JERK

    # 加载模型（头部HDV的自由驾驶模型、尾部HDV的follow模型）
    self.hdvH_tf_model_pred = load_model(TRANSFORMER_FREE_DRIVING_MODEL, custom_objects={'PositionalEncoding': PositionalEncoding})
    self.hdvF_tf_model_true = load_model(TRANSFORMER_FOLLOW_MODEL_TRUE, custom_objects={'PositionalEncoding': PositionalEncoding})
    self.hdvF_tf_model_pred = load_model(TRANSFORMER_FOLLOW_MODEL_PRED, custom_objects={'PositionalEncoding': PositionalEncoding})

  def reset(self, hdvH_data, hdvF_data):
    # 时间步
    self.time_step = 0

    # 总帧数
    self.total_step = 30 # todo 换道时长取3秒

    # 记录HDV初始状态数据
    self.hdvH_data= np.array(hdvH_data)
    self.hdvF_data= np.array(hdvF_data)
    self._create_n_past_data()

    # 获取初始状态
    self.state, self.kalman_state, self.control_state, self.predict_state = self._get_init_state() # 只用最后一帧数据初始化环境状态
    print("【init】 state: ", np.array(self.state).tolist())
    self.last_state = self.state
    self.last_kalman_state = self.kalman_state
    self.last_actions = [0, 0, 0, 0, 0, 0, 0, 0, 0] # 上一时刻的动作

    normalized_state = self._normalization_state(self.state)
    return normalized_state, self.state, self.kalman_state, self.control_state
  def step(self, actions):
    # 更新时间步
    self.time_step += 1
    print("===== Step ", self.time_step, " =====")
    
    # 给动作添加控制扰动
    self.original_actions = actions
    actions = self._add_control_noise(actions)

    # 格式化动作序列
    cavM_accs, cavM_deltas, cav1_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs = self._format_actions(actions)
    hdvH_accs, hdvF_accs = self._get_hdv_cur_accs()

    # 迭代环境状态
    self._run(hdvH_accs, cavM_accs, cavM_deltas, cav1_accs, hdvF_accs)

    # 更新历史记录
    self._update_n_past_data(hdvH_accs, cavM_accs, cav1_accs, hdvF_accs)

    # 获取预测控制行为
    hdvH_predict_accs, hdvF_predict_accs, hdvH_predict_stds, hdvF_predict_stds = self._get_hdv_predict_accs()

    # 获取状态
    self.state, self.kalman_state, self.control_state, self.predict_state = self._get_state(cavM_accs, cavM_deltas, hdvH_accs, cav1_accs, hdvF_accs, hdvH_predict_accs, hdvF_predict_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs, hdvH_predict_stds, hdvF_predict_stds)
    normalized_state = self._normalization_state(self.state)

    # 计算奖励
    reward, reward_info = self._get_reward(cavM_accs, cavM_deltas, cav1_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs, hdvH_predict_accs, hdvF_predict_accs)
    print("reward: ", reward)
    print("reward_acc: ", reward_info["reward_acc"])
    print("reward_delta: ", reward_info["reward_delta"])
    print("reward_v_diff: ", reward_info["reward_v_diff"])
    print("reward_y_diff: ", reward_info["reward_y_diff"])
    print("reward_acc_jerk: ", reward_info["reward_acc_jerk"])
    print("reward_delta_jerk: ", reward_info["reward_delta_jerk"])
    print("reward_psi: ", reward_info["reward_psi"])
    print("reward_v_x: ", reward_info["reward_v_x"])
    print("reward_x", reward_info["reward_x"])

    # 判断是否结束
    done = self._is_done()

    # 记录环境状态
    self.last_state = copy.deepcopy(self.state)
    self.last_actions = copy.deepcopy(actions)
    return normalized_state, reward, done, reward_info, self.state, self.kalman_state, self.control_state
  
  def _create_n_past_data(self):
    # hdvF的历史数据
    self.hdvF_past_data = self.hdvF_data[:self.hdvF_n_past, :].copy()

    # cav1的历史数据(用真实的前车数据进行初始化，因为随意设置的数据可能导致不符合动力学模型，最终导致transformer模型的预测很离谱)
    self.cav1_past_data = self.hdvH_data[:self.hdvF_n_past, :].copy()

    # 修正hdvH_data的y坐标数据
    y_diffs = self.hdvH_data[:, 0] - self.hdvF_data[:, 0]
    self.hdvH_data[:, 0] = self.hdvH_data[:, 0] + y_diffs[:]

    # 创建hdvH的历史数据
    self.hdvH_past_data = self.hdvH_data[:self.hdvH_n_past, :].copy()

    # 创建cavM的历史数据(意义不大，其实只需要生成最后一帧)
    self.cavM_past_data = np.zeros((self.hdvH_n_past, 4))
    self.cavM_past_data[:, 0] = self.hdvH_past_data[:, 0] - (y_diffs[:self.hdvH_n_past] / 2)
    self.cavM_past_data[:, 1] = 0
    self.cavM_past_data[:, 2] = self.hdvH_past_data[:, 2] + random.uniform(-2, 2)
    self.cavM_past_data[:, 3] = self.hdvH_past_data[:, 3] + random.uniform(-0.3, 0.3)

  def _update_n_past_data(self, hdvH_accs, cavM_accs, cav1_accs, hdvF_accs):
    # 获取滤波数据
    hdvH_y, _, hdvH_v, hdvH_covariances = self.hdvH_model.get_filter_state()
    cavM_y, cavM_x, cavM_psi, cavM_v, cavM_v_y, cavM_v_x, cavM_covariances = self.cavM_model.get_filter_state()
    cav1_y, _, cav1_v, cav1_covariances = self.cav1_model.get_filter_state()
    hdvF_y, _, hdvF_v, hdvF_covariances = self.hdvF_model.get_filter_state()

    # 更新hdvH的历史数据
    self.hdvH_past_data = np.delete(self.hdvH_past_data, 0, axis=0) # 删除第一行
    self.hdvH_past_data = np.append(self.hdvH_past_data, np.array([[hdvH_y, self.hdvH_past_data[-1, 1], hdvH_v, hdvH_accs[0]]]), axis=0) # 添加新行

    # 更新cavM的历史数据
    self.cavM_past_data = np.delete(self.cavM_past_data, 0, axis=0) # 删除第一行
    self.cavM_past_data = np.append(self.cavM_past_data, np.array([[cavM_y, cavM_x, cavM_v, cavM_accs[0]]]), axis=0) # 添加新行

    # 更新cav1的历史数据
    self.cav1_past_data = np.delete(self.cav1_past_data, 0, axis=0) # 删除第一行
    self.cav1_past_data = np.append(self.cav1_past_data, np.array([[cav1_y, self.hdvH_past_data[-1, 1], cav1_v, cav1_accs[0]]]), axis=0) # 添加新行

    # 更新hdvF的历史数据
    self.hdvF_past_data = np.delete(self.hdvF_past_data, 0, axis=0) # 删除第一行
    self.hdvF_past_data = np.append(self.hdvF_past_data, np.array([[hdvF_y, self.hdvH_past_data[-1, 1], hdvF_v, hdvF_accs[0]]]), axis=0) # 添加新行

  def _get_hdv_cur_accs(self):
    """
      描述：获取HDV当前时刻的加速度
    """
    # 获取hdvH的第一帧加速度
    hdvH_shift = self.time_step + self.hdvH_n_past - 1
    hdvH_acc = self.hdvH_data[hdvH_shift][3] # 纵向加速度

    # 获取hdvF的第一帧加速度
    hdvF_accs, hdvF_stds = self._get_hdvF_predict_accs(self.hdvF_tf_model_true)

    return [hdvH_acc], [hdvF_accs[0]]
  
  def _get_hdv_predict_accs(self):
    """
      描述：获取hdv的预测加速度
    """
    # 获取hdvH的预测加速度序列
    hdvH_accs, hdvH_stds = self._get_hdvH_predict_accs(self.hdvH_tf_model_pred)

    # 获取hdvF的预测加速度序列
    hdvF_accs, hdvF_stds = self._get_hdvF_predict_accs(self.hdvF_tf_model_pred)

    return hdvH_accs, hdvF_accs, hdvH_stds, hdvF_stds
  
  def _get_hdvF_predict_accs(self, model):
    """
      描述：获取hdvF的预测加速度
    """
    # 生成预测需要的历史数据
    hdvF_input_data = []
    for i in range(self.hdvF_n_past):
      temp = []
      temp.append(self.cav1_past_data[i][2]) # 纵向速度
      temp.append(self.hdvF_past_data[i][2]) # 纵向速度
      temp.append(self.cav1_past_data[i][3]) # 纵向加速度
      temp.append(self.hdvF_past_data[i][3]) # 纵向加速度
      temp.append(self.cav1_past_data[i][0] - self.hdvF_past_data[i][0]) # 车头间距
      hdvF_input_data.append(temp)
    hdvF_input_data = np.array(hdvF_input_data)
    # print("hdvF_input_data: ", len(hdvF_input_data), len(hdvF_input_data[0])) # 30 5

    # 归一化
    hdvF_input_data_normalized, min_max_list = normalization([hdvF_input_data], [0, 1, 2, 3, 4], [True, True, False, False, True])

    # 生成预测数据
    n_samples = 5
    mc_preds = np.stack([model(hdvF_input_data_normalized, training=True) for _ in range(n_samples)])
    # print("mc_preds: ", len(mc_preds), len(mc_preds[0]), len(mc_preds[0][0])) # 1 10 1

    # 反归一化
    mc_preds = np.stack([anit_normalization(mc_preds[i], [0], [min_max_list[3]], [False]) for i in range(n_samples)])
    # print("hdvF_predict_datas(anit): ", hdvF_predict_datas[0])

    # 取均值
    hdvF_predict_datas = mc_preds.mean(axis=0)

    # 计算方差(量化不确定性)
    hdvF_predict_datas_std = mc_preds.std(axis=0)

    return hdvF_predict_datas[0][:, 0], hdvF_predict_datas_std[0][:, 0]
  
  def _get_hdvH_predict_accs(self, model):
    """
      描述：获取hdvH的预测加速度
    """
    # 生成预测需要的历史数据
    hdvH_input_data = []
    for i in range(self.hdvH_n_past):
      temp = []
      temp.append(self.hdvH_past_data[i][2]) # 纵向速度
      temp.append(self.hdvH_past_data[i][3]) # 纵向加速度
      hdvH_input_data.append(temp)
    hdvH_input_data = np.array(hdvH_input_data)
    # print("hdvH_input_data: ", len(hdvH_input_data), len(hdvH_input_data[0])) # 30 2

    # 归一化
    hdvH_input_data_normalized, min_max_list = normalization([hdvH_input_data], [0, 1], [True, False])

    # 生成预测数据
    n_samples = 5
    mc_preds = np.stack([model(hdvH_input_data_normalized, training=True) for _ in range(n_samples)])

    # 反归一化
    mc_preds = np.stack([anit_normalization(mc_preds[i], [0], [min_max_list[1]], [False]) for i in range(n_samples)])

    # 取均值
    hdvH_predict_datas = mc_preds.mean(axis=0)
    # print("hdvH_predict_datas: ", len(hdvH_predict_datas), len(hdvH_predict_datas[0]), len(hdvH_predict_datas[0][0])) # 1 10 1

    # 计算方差(量化不确定性)
    hdvH_predict_datas_std = mc_preds.std(axis=0)
    # print("hdvH_predict_datas_std: ", len(hdvH_predict_datas_std), len(hdvH_predict_datas_std[0]), len(hdvH_predict_datas_std[0][0])) # 1 10 1

    return hdvH_predict_datas[0][:, 0], hdvH_predict_datas_std[0][:, 0]

  def _run(self, hdvH_cur_accs, cavM_accs, cavM_deltas, cav1_accs, hdvF_accs):
    """
      描述：迭代HDV和CAV的状态
    """
    # 更新头部HDV的状态
    self.hdvH_model.update(hdvH_cur_accs[0])

    # 更新换道CAV的状态
    self.cavM_model.update(cavM_accs[0], cavM_deltas[0])

    # 更新cav1的状态
    self.cav1_model.update(cav1_accs[0])

    # 更新尾部HDV的状态
    self.hdvF_model.update(hdvF_accs[0])

  def _is_done(self):
    """
      描述：判断是否结束
    """
    cavM_x = self.state[1]
    
    return self.time_step >= self.total_step or cavM_x >= self.lane_width

  def _get_init_state(self):
    """
      描述：获取初始状态
    """
    # 获取当前状态（最后一帧充当当前状态）
    hdvH_cur_state = self.hdvH_past_data[-1, :]
    cavM_cur_state = self.cavM_past_data[-1, :]
    cav1_cur_state = self.cav1_past_data[-1, :]
    hdvF_cur_state = self.hdvF_past_data[-1, :]

    # 创建动态模型
    self.hdvH_model = PointMassModel(self.dt, hdvH_cur_state[0], hdvH_cur_state[1], hdvH_cur_state[2])
    self.cavM_model = BicycleModel(self.dt, 1, cavM_cur_state[0], 0, cavM_cur_state[2], math.pi / 2)
    self.cav1_model = PointMassModel(self.dt, cav1_cur_state[0], cav1_cur_state[1], cav1_cur_state[2])
    self.hdvF_model = PointMassModel(self.dt, hdvF_cur_state[0], hdvF_cur_state[1], hdvF_cur_state[2])

    return self._get_state([0], [0], [0], [0], [0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0])

  def _get_state(self, cavM_accs, cavM_deltas, hdvH_accs, cav1_accs, hdvF_accs, hdvH_predict_accs, hdvF_predict_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs, hdvH_predict_stds, hdvF_predict_stds):
    """
      描述：生成每一次迭代的环境状态
    """
    # 获取滤波数据
    hdvH_y, _, hdvH_v, hdvH_covariances = self.hdvH_model.get_filter_state()
    cavM_y, cavM_x, cavM_psi, cavM_v, cavM_v_y, cavM_v_x, cavM_covariances = self.cavM_model.get_filter_state()
    cav1_y, _, cav1_v, cav1_covariances = self.cav1_model.get_filter_state()
    hdvF_y, _, hdvF_v, hdvF_covariances = self.hdvF_model.get_filter_state()

    # 获取真实数据
    hdvH_true_y, _, hdvH_true_v = self.hdvH_model.get_state()
    cavM_true_y, cavM_true_x, cavM_true_psi, cavM_true_v, cavM_true_v_y, cavM_true_v_x = self.cavM_model.get_state()
    cav1_true_y, _, cav1_true_v = self.cav1_model.get_state()
    hdvF_true_y, _, hdvF_true_v = self.hdvF_model.get_state()

    # 获取测量数据
    hdvH_mea_y, _, hdvH_mea_v = self.hdvH_model.get_measure_state()
    cavM_mea_y, cavM_mea_x, cavM_mea_psi, cavM_mea_v, cavM_mea_v_y, cavM_mea_v_x = self.cavM_model.get_measure_state()
    cav1_mea_y, _, cav1_mea_v = self.cav1_model.get_measure_state()
    hdvF_mea_y, _, hdvF_mea_v = self.hdvF_model.get_measure_state()

    # 获取预测数据
    hdvH_predict_states = self.hdvH_model.calc_predict_state(np.array(hdvH_predict_accs)[0: 2])
    cavM_predict_states = self.cavM_model.calc_predict_state(cavM_predict_accs, cavM_predict_deltas)
    cav1_predict_states = self.cav1_model.calc_predict_state(cav1_predict_accs)
    hdvF_predict_states = self.hdvF_model.calc_predict_state(np.array(hdvF_predict_accs)[0: 2])

    # 当前时刻真实的状态（这里cavM_x可以用来表示cavM换道的横向累计距离）
    original_state = [
                      hdvH_y - cavM_y, cavM_x, cavM_v, cavM_psi, hdvH_v,
                      cavM_y - cav1_y, cav1_y - hdvF_y, cav1_v, hdvF_v, hdvH_accs[0], hdvF_accs[0],
                      cavM_covariances[1], cavM_covariances[0], cavM_covariances[2], cavM_covariances[3],
                      hdvH_predict_accs[0], hdvH_predict_accs[1], # hdvH的两帧预测加速度
                      hdvF_predict_accs[0], hdvF_predict_accs[1], # hdvF的两帧预测加速度
                      hdvH_predict_stds[0], hdvH_predict_stds[1], # hdvH的两帧预测方差
                      hdvF_predict_stds[0], hdvF_predict_stds[1], # hdvF的两帧预测方差
                     ] # 状态空间需要的信息
    # original_other_state = [cavM_v_y, cavM_v_x, cavM_accs[0], cavM_deltas[0], hdvH_y, cavM_y, cav1_accs[0], cav1_y, hdvF_y] # 其他计算时需要的信息
    original_control_state = {
      "hdvH_accs": hdvH_accs,
      "cavM_accs": cavM_accs,
      "cavM_deltas": cavM_deltas,
      "cav1_accs": cav1_accs,
      "hdvF_accs": hdvF_accs,
    }
    orginal_kalman_state = {
      "hdvH_true_state": [hdvH_true_y, hdvH_true_v],
      "cavM_true_state": [cavM_true_y, cavM_true_x, cavM_true_psi, cavM_true_v, cavM_true_v_y, cavM_true_v_x],
      "cav1_true_state": [cav1_true_y, cav1_true_v],
      "hdvF_true_state": [hdvF_true_y, hdvF_true_v],
      "hdvH_mea_state": [hdvH_mea_y, hdvH_mea_v],
      "cavM_mea_state": [cavM_mea_y, cavM_mea_x, cavM_mea_psi, cavM_mea_v, cavM_mea_v_y, cavM_mea_v_x],
      "cav1_mea_state": [cav1_mea_y, cav1_mea_v],
      "hdvF_mea_state": [hdvF_mea_y, hdvF_mea_v],
      "hdvH_filter_state": [hdvH_y, hdvH_v],
      "hdvH_filter_covariances": hdvH_covariances,
      "cavM_filter_state": [cavM_y, cavM_x, cavM_psi, cavM_v, cavM_v_y, cavM_v_x],
      "cavM_filter_covariances": cavM_covariances,
      "cav1_filter_state": [cav1_y, cav1_v],
      "cav1_filter_covariances": cav1_covariances,
      "hdvF_filter_state": [hdvF_y, hdvF_v],
      "hdvF_filter_covariances": hdvF_covariances,
    }
    original_predict_state = {
      "hdvH_predict_states": hdvH_predict_states,
      "cavM_predict_states": cavM_predict_states,
      "cav1_predict_states": cav1_predict_states,
      "hdvF_predict_states": hdvF_predict_states,
    }
    # print("hdvH_true_state: ", np.array(orginal_kalman_state["hdvH_true_state"]).tolist())
    # print("hdvH_mea_state: ", np.array(orginal_kalman_state["hdvH_mea_state"]).tolist())
    # print("hdvH_filter_state: ", np.array(orginal_kalman_state["hdvH_filter_state"]).tolist())
    print("cavM_true_state: ", np.array(orginal_kalman_state["cavM_true_state"]).tolist())
    print("cavM_mea_state: ", np.array(orginal_kalman_state["cavM_mea_state"]).tolist())
    print("cavM_filter_state: ", np.array(orginal_kalman_state["cavM_filter_state"]).tolist())
    # print("cav1_true_state: ", np.array(orginal_kalman_state["cav1_true_state"]).tolist())
    # print("cav1_mea_state: ", np.array(orginal_kalman_state["cav1_mea_state"]).tolist())
    # print("cav1_filter_state: ", np.array([cav1_y, cav1_v]).tolist())
    # print("hdvF_true_state: ", np.array(orginal_kalman_state["hdvF_true_state"]).tolist())
    # print("hdvF_mea_state: ", np.array(orginal_kalman_state["hdvF_mea_state"]).tolist())
    # print("hdvF_filter_state: ", np.array(orginal_kalman_state["hdvF_filter_state"]).tolist())
    return original_state, orginal_kalman_state, original_control_state, original_predict_state

  def _get_reward(self, cavM_accs, cavM_deltas, cav1_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs, hdvH_predict_accs, hdvF_predict_accs):
    """
      描述：计算奖励
    """
    # !目标函数
    # 加速度
    reward_acc, _ = r_acc(cavM_accs, cav1_accs, cavM_predict_accs, cav1_predict_accs)

    # 转向角
    reward_delta, _ = r_delta(cavM_deltas, cavM_predict_deltas)

    # 纵向速度差
    reward_v_diff, _ = r_v_diff(self.state, self.kalman_state, self.predict_state)

    # 纵向间距
    reward_y_diff, _ = r_y_diff(self.state, self.kalman_state, self.thdv, self.tcav, self.predict_state)

    # !约束
    # 加速度变化率
    last_cavM_accs, last_cavM_deltas, last_cav1_accs, last_cavM_predict_accs, last_cavM_predict_deltas, last_cav1_predict_accs = self._format_actions(self.last_actions)
    reward_acc_jerk, _ = r_acc_jerk(self.dt, cavM_accs, last_cavM_accs, cav1_accs, last_cav1_accs, cavM_predict_accs, cav1_predict_accs)

    # 转向角变化率
    reward_delta_jerk, _ = r_delta_jerk(self.dt, cavM_deltas, last_cavM_deltas, cavM_predict_deltas)

    # 航向角变化率
    reward_psi, _ = r_psi(self.time_step, self.state, self.last_state, self.predict_state)

    # 横向速度
    reward_v_x, _ = r_v_x(self.last_kalman_state, self.kalman_state, self.predict_state, self.time_step)

    # 横向间距
    reward_x, _ = r_x(self.last_kalman_state, self.kalman_state, self.predict_state, self.lane_width, self.time_step)

    # todo 将transformer的预测方差引入奖励函数中。计划是基于方差来区分策略（如果方差大, 则需要执行保守策略）

    # todo 碰撞惩罚

    # !汇总奖励
    reward_acc *=  0.5
    reward_delta *= 1
    reward_v_diff *= 1
    reward_y_diff *= 1.5
    reward_acc_jerk *= 0.5
    reward_delta_jerk *= 1.5
    reward_psi *= 2
    reward_v_x *= 2
    reward_x *= 5
    reward = reward_acc + reward_delta + reward_v_diff + reward_y_diff + reward_acc_jerk + reward_delta_jerk + reward_psi + reward_v_x + reward_x
  
    reward_info = {
      "reward_acc": reward_acc,
      "reward_delta": reward_delta,
      "reward_v_diff": reward_v_diff,
      "reward_y_diff": reward_y_diff,
      "reward_acc_jerk": reward_acc_jerk,
      "reward_delta_jerk": reward_delta_jerk,
      "reward_psi": reward_psi,
      "reward_v_x": reward_v_x,
      "reward_x": reward_x,
    }
    return reward, reward_info

  def _normalization_state(self, state):
    """
      描述：对状态做归一化
    """
    copy_state = copy.deepcopy(state)

    copy_state[0] /= 40 # 纵向间距归一化
    copy_state[1] /= self.lane_width # 横向间距归一化
    copy_state[2] /= 30 # cavM的轴向速度归一化
    copy_state[3] /= math.pi / 2  # 航向角归一化
    copy_state[4] /= 30 # hdvH的纵向速度归一化
    copy_state[5] /= 40 # 纵向间距归一化
    copy_state[6] /= 40 # 纵向间距归一化
    copy_state[7] /= 30 # cav1的纵向速度归一化
    copy_state[8] /= 30 # hdvF的纵向速度归一
    copy_state[9] /= self.acc_bound # hdvH的加速度归一化
    copy_state[10] /= self.acc_bound # hdvF的加速度归一化
    copy_state[11] /= 1 # cavM的y协方差
    copy_state[12] /= 1 # cavM的x协方差
    copy_state[13] /= 1 # cavM的v协方差
    copy_state[14] /= 1 # cavM的psi协方差

    return np.array(copy_state)

  def _format_actions(self, actions):
    """
      描述：格式化本次生成的动作
    """
    # 换道cav的action（当前执行和预测）
    cavM_accs = []
    cavM_deltas = []
    cavM_predict_accs = []
    cavM_predict_deltas = []
    # 跟随cav的action（当前执行和预测）
    cav1_accs = []
    cav1_predict_accs = []

    # 将原始的动作结合各自的bound生成对应的动作
    cavM_accs.append(actions[0] * ACC_BOUND)
    cavM_predict_accs.append(actions[1] * ACC_BOUND)
    cavM_predict_accs.append(actions[2] * ACC_BOUND)
    cavM_deltas.append(actions[3] * DELTA_BOUND)
    cavM_predict_deltas.append(actions[4] * DELTA_BOUND)
    cavM_predict_deltas.append(actions[5] * DELTA_BOUND)
    cav1_accs.append(actions[6] * ACC_BOUND)
    cav1_predict_accs.append(actions[7] * ACC_BOUND)
    cav1_predict_accs.append(actions[8] * ACC_BOUND)

    return cavM_accs, cavM_deltas, cav1_accs, cavM_predict_accs, cavM_predict_deltas, cav1_predict_accs
  
  def _add_control_noise(self, actions):
    new_actions = copy.deepcopy(actions)

    # 加速度
    new_actions[0] = (actions[0] * ACC_BOUND + np.random.normal(0, ACC_NOISE_STD)) / ACC_BOUND

    # 转向角
    new_actions[3] = (actions[3] * DELTA_BOUND + np.random.normal(0, DELTA_NOISE_STD)) / DELTA_BOUND

    # 跟随cav的加速度
    new_actions[6] = (actions[6] * ACC_BOUND + np.random.normal(0, ACC_NOISE_STD)) / ACC_BOUND

    return new_actions
